package org.erp.common.aspect;

import lombok.RequiredArgsConstructor;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.erp.common.annotation.StockLockBatch;
import org.erp.common.context.StockLockContext;
import org.erp.common.manager.StockBatchLockManager;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.UUID;
import java.util.function.Consumer;

@Aspect
@Component
@RequiredArgsConstructor
public class StockBatchLockAspect {

    private final StringRedisTemplate redis;

    private final StockBatchLockManager stockBatchLockManager;

    @Around("@annotation(stockLockBatch)")
    public Object lockBatch(ProceedingJoinPoint pjp, StockLockBatch stockLockBatch) throws Throwable {

        Object arg = pjp.getArgs()[0];

        // ======================
        // directIds 分支：参数本身是 List<Long> 或 List<String>
        // ======================
        if (stockLockBatch.directIds()) {
            if (!(arg instanceof List<?> idsList)) {
                throw new IllegalArgumentException("directIds=true 时，方法参数必须是 List");
            }

            // 转为字符串 ID
            List<String> allIds = idsList.stream()
                    .map(String::valueOf)
                    .toList();

            return processIds(pjp, allIds, stockLockBatch, null);
        }

        // ======================
        // DTO 分支（原逻辑）
        // ======================
        Object dto = arg;

        Field listField = dto.getClass().getDeclaredField(stockLockBatch.listField());
        listField.setAccessible(true);
        List<?> list = (List<?>) listField.get(dto);

        if (list == null || list.isEmpty()) {
            return pjp.proceed();
        }

        Object first = list.get(0);
        if (first instanceof Long || first instanceof String) {

            List<String> allIds = list.stream()
                    .map(String::valueOf)
                    .toList();

            return processIds(pjp, allIds, stockLockBatch, (freeIds) -> {
                List<?> filtered = list.stream()
                        .filter(item -> freeIds.contains(String.valueOf(item)))
                        .toList();
                try {
                    listField.set(dto, filtered);
                } catch (IllegalAccessException e) {
                    throw new RuntimeException(e);
                }
            });
        }

        Field idField = list.get(0).getClass().getDeclaredField(stockLockBatch.idField());
        idField.setAccessible(true);

        List<String> allIds = new ArrayList<>();
        for (Object item : list) {
            Object idVal = idField.get(item);
            allIds.add(String.valueOf(idVal));
        }

        return processIds(pjp, allIds, stockLockBatch, (freeIds) -> {
            // 回写 DTO，只保留可用的对象
            List<?> filtered = list.stream()
                    .filter(item -> {
                        try {
                            return freeIds.contains(String.valueOf(idField.get(item)));
                        } catch (Exception e) {
                            return false;
                        }
                    })
                    .toList();
            try {
                listField.set(dto, filtered);
            } catch (IllegalAccessException e) {
                throw new RuntimeException(e);
            }
        });
    }

    private Object processIds(ProceedingJoinPoint pjp,
                              List<String> allIds,
                              StockLockBatch stockLockBatch,
                              Consumer<List<String>> dtoCallback) throws Throwable {

        String busyKey = stockLockBatch.busyKey();

        // 过滤已占用
        List<String> lockedIds = redis.opsForSet().intersect(busyKey, new HashSet<>(allIds))
                .stream().toList();

        List<String> freeIds = allIds.stream()
                .filter(id -> !lockedIds.contains(id))
                .toList();

        if (dtoCallback != null) {
            dtoCallback.accept(freeIds);
        }

        if (freeIds.isEmpty()) {
            return pjp.proceed();
        }

        // 生成批次锁
        String batchKey = "lock:batch:" + UUID.randomUUID();

        StockLockContext.setBatchKey(batchKey);

        String script = """
                for i=1,#KEYS do
                    redis.call('sadd', ARGV[1], KEYS[i])
                end
                redis.call('set', ARGV[2], table.concat(KEYS, ','), 'EX', ARGV[3])
                return 1
                """;

        redis.execute(
                RedisScript.of(script, Long.class),
                freeIds, busyKey, batchKey, String.valueOf(stockLockBatch.ttlSeconds())
        );

        try {
            return pjp.proceed();
        } finally {
            try {
                stockBatchLockManager.partialUnlockIfNecessary(busyKey);
            } finally {
                StockLockContext.clear();
            }
        }
    }
}
